"""

    File: rtmp.py
    Description: 
    
        This module implements Adobe's proprietary RTMP protocol.
        
    Author: Kyle Vogt
    Date  : October 8th, 2007
    Copyright (c) 2007, Justin.tv, Inc.
    
"""

from twisted.internet import reactor, protocol, defer
from twisted.python import log
from pms.rtmp.encoder_py import Encoder
from pms.rtmp.decoder_py import Decoder
from pms.rtmp import amf
from pms.rtmp.netstream import NetStream
from pms.rtmp.framedropper import FrameDropper
import time, traceback, struct

class Base(protocol.Protocol):

    def __init__(self):
        self.handshaked = False
        self.connected = False
        self.state = False
        self.running = True
        self.handshake_data = ''
        self.buffer = []
        self.decoder = Decoder()
        self.encoder = Encoder()
        self.dropper = FrameDropper()
        self.datas = []
        self.amf = amf.Amf()
        self.streams = {}
        self.count = 0
        self.bytesRead = 0
        self.bytesSent = 0
        self.last_data_time = time.time()
        # Stream starts off with no access code
        self.password = None
        self.prioritize = True
        self.tcUrl = self.pageUrl = self.swfUrl = self.app = None
        
        #self.byte_log = open('data.log', 'w')
        
    def create_netstream(self, streamId):
        try:
            self.streams[int(streamId)] = NetStream(self, int(streamId))
        except:
            log.msg('Could not create netstream %s' % streamId)
            traceback.print_exc()
        
    def run_state(self):
        if self.state:
            try:
                self.state()
            except:
                log.msg("Error in state: %s" % self.state)
                traceback.print_exc()
        
    def handleEvents(self):
        "Process commands sent over this NetConnection"
        while len(self.decoder.events):
            event = self.decoder.events.pop(0)
            data, streamId = event['data'], int(event['id'])
            if not self.factory.clientMode and streamId > 0 and streamId not in self.streams:
                log.msg("Bad NetStream id: %s" % streamId)
                break
            else: 
                if streamId not in self.streams: 
                    self.create_netstream(streamId)
                stream = self.streams[streamId]
            if data and isinstance(data[0], str):                 
                command = data[0].replace('@', '')
                if command == '_result': 
                    self.__result(data, stream)
                    continue
                func = None
                try:
                    func = getattr(self, '_' + command)
                except AttributeError:
                    #log.msg("No handler for command %s, sending to plugins" % command)
                    try:
                        self.factory.core.plug('invoke', data, stream)
                    except:
                        log.msg('Error calling invoke on plugins!')
                        traceback.print_exc()
                try:
                    if func: func(data, stream)
                except:
                    log.msg("Error calling %s" % command)
                    traceback.print_exc()
        for stream in self.streams.values():
            packets = [packet for packet in self.decoder.packets if packet['id'] == stream.streamId]
            if stream.stream and len(packets):
                stream.stream.addPackets(packets)
        #log.msg('Types: %s Streams: %s Objects: %s' % (', '.join([str(packet['type']) for packet in self.decoder.packets]), ', '.join([str(packet['id']) for packet in self.decoder.packets]), ', '.join([str(packet['object']) for packet in self.decoder.packets])))
        self.decoder.packets = []
            
    def _connect(self, data, stream):    
        self.connected = True
        if len(data) > 2:
            for name in ['app', 'tcUrl', 'pageUrl', 'swfUrl']:
                setattr(self, name, data[2].get(name, None))
        # Report connection success
        obj = {    'level' : 'status', 
                'code' : 'NetConnection.Connect.Success',
                   'description' : 'Connection succeeded'}
        self.invoke('_result', 1.0, obj)
        # Set chunk size and reconfigure encoder
        self.send(amf.META_CHANNEL, amf.CHUNK_SIZE, 0, struct.pack('!i', 65535))
        self.encoder.setChunkSize(65535)
        
    def _createStream(self, data, stream):
        id = data[1]
        log.msg("Created netstream %s" % id)
        self.streams[id] = NetStream(self, id)
        # client stream, server stream 
        self.invoke('_result', id, id)
        
    def _closeStream(self, data, stream):
        stream.close()
        log.msg("Closed stream %s" % stream.streamId)

    def _releaseStream(self, data, stream):
        pass
    
    def _deleteStream(self, data, stream):
        id = data[3]
        if id in self.streams and not self.streams[id].stream:
            self.streams[id].close()
            del self.streams[id]
            log.msg("Deleted stream %s" % id)
                                
    def __result(self, data, stream):
        index = len(data) - 1
        try:
            level = data[index].get('level')
            code = data[index].get('code')
            description = data[index].get('description')
        except:
            code = "NetStream created with id %s" % data[3]
            self.factory.createStreamDeferred.callback(data)
            #self.factory.createStreamDeferred = defer.Deferred()                                
        # Handle the results
        if code == "NetConnection.Connect.Success":
            log.msg('Connect success!')
            self.remoteConnected = True
            self.factory.connectDeferred.callback(data)
        elif code == "NetStream.Publish.Start":
            self.remotePublishing = True
            self.factory.publishDeferred.callback(data)
        elif code == "NetStream.Play.Start":
            if not self.remotePlaying:
                self.factory.playDeferred.callback(data)
            self.remotePlaying = True
        elif code == "NetStream.Play.StreamNotFound":
            self.remotePlaying = False
            self.factory.playDeferred.errback(data)
        else:
            log.msg("Unhandled code: %s" % code)
            
    def _onStatus(self, data, stream):
        self.__result(data, stream)
        
    def _setDataFrame(self, data, stream):
        if data[1] == 'onMetaData':
            stream.stream.meta = data[2]
        
    def _ping(self, data, stream):
        self.invoke('pong', 0.0, 'Server connection is alive.')

    def setState(self, state):
        try: 
            self.state = getattr(self, state)
            log.msg("State is %s" % state)
        except:
            log.msg("No handler for state: %s (looked for %s)" % (state, state))
            self.shutdown()
        self.run_state()
            
    def send(self, channel, type, stream, data, timecode=None):
        if not timecode: timecode = self.timecode()
        data = self.encoder.process(data, channel, stream, type, timecode, True)
        self.write(data)    
        
    def invoke(self, command, index, objs=None, streamId=0, channel=None, insertNone=True):
        "Invoke a server side function on a remote server"
        args = [command, index]
        if insertNone:
            args.append(None)
        if objs: 
            if not isinstance(objs, list): objs = [objs]
            args.extend(objs)
        if not channel: 
            channel = amf.CONNECTION_CHANNEL
        #log.msg('sending args: %s' % args)
        self.send(channel, amf.INVOKE, streamId, self.amf.encode(args))
        
    def ping(self, data):
        "Send a ping to the server"
        self.send(amf.META_CHANNEL, amf.PING, 0, data, timecode=0)
        
    def server_bw(self, data):
        self.send(amf.META_CHANNEL, amf.SERVER_BW, 0, data)
        
    def client_bw(self, data):
        self.send(amf.META_CHANNEL, amf.CLIENT_BW, 0, data)
        
    def dataReceived(self, data):
        "Read the raw data"
        self.last_data_time = time.time()
        if not self.running: return
        bytes = len(data)
        self.bytesRead += bytes
        self.factory.global_bytes_read += bytes
        self.count += bytes
        if self.count > 524288:
            # It's time to ping the client
            self.send(amf.META_CHANNEL, amf.BYTES_READ, 0, struct.pack("!i", self.bytesRead % 2147483648))
            self.count -= 524288
        if not self.handshaked: 
            # We're handshaking, so buffer the data
            self.handshake_data += data
            self.run_state()
        elif self.factory.clientMode:
            self.decoder.process(data)
            self.run_state()
        else:    
            self.buffer.append(data)
            
    def poll(self):
        "Called by the scheduler"
        try:
            self.decoder.process(''.join(self.buffer))
        except:
            log.msg('Error - Unable to process buffer during poll()')
            traceback.print_exc()
        self.buffer = []
        self.run_state()
            
    def write(self, data):
        "Wrapper for writing to the socket"
        self.transport.write(data)
        bytes = len(data)
        self.bytesSent += bytes
        self.factory.global_bytes_sent += bytes
        
    def timecode(self):
        "Return the current timecode in ms"
        return int((time.time() - self.connectionTime) * 1000.0)
        
    def connectionMade(self):
        "Handle a TCP connection"
        self.connectionTime = time.time()
        self.ip = self.transport.getPeer().host
        self.port = self.transport.getPeer().port
        log.msg("Connection at %s:%s" % (self.ip, self.port))
        # Create default netstream on this connection
        self.streams[0] = NetStream(self, 0)
        # Check connection limits
        if self.factory.addClient(self):
            self.setState('handshake')
        else:
            log.msg('Cannot accept client %s, at connection limit' % self.ip)
            self.shutdown()
            
    def connectionLost(self, reason):
        "If connection breaks, call shutdown procedure"
        log.msg("%s disconnected" % self.ip)
        self.shutdown(False)    
        self.factory.removeClient(self)
            
    def shutdown(self, drop=True):
        "Shut down processing and update state information"
        self.running = False
        self.connected = False
        for stream in self.streams.values(): stream.close()
        self.transport.dataBuffer = ''
        self.transport._tempDataBuffer = [] # will be added to dataBuffer in doWrite
        self.transport._tempDataLen = 0
        self.dropper.stopProducing()
        # Give the garbage collector some help!
        self.streams = {}
        # Lose connection if we haven't already
        if drop: 
            self.transport.loseConnection()

class NetConnection(Base):

    def __init__(self):
        Base.__init__(self)
        self.publishing = False
        self.playing = False
            
    def readClientHandshake(self):
        if len(self.handshake_data) >= 1537:
            shake, self.handshake_data = self.handshake_data[1:1537], self.handshake_data[1537:]
            #reply = chr(3) + (chr(0) * 1536) + shake
            reply = chr(3) + '\x02\xf7J-\x00\x00\x00\x00\x92\x00\xbb\x00x\x00\xd9\x00.\x00\x07\x004\x00\xc5\x00\n\x00\x93\x000\x00\xf1\x00&\x00_\x00l\x00]\x00\x82\x00k\x00\xe8\x00\t\x00\x1e\x00\xb7\x00\xa4\x00\xf5\x00\xfa\x00C\x00\xa0\x00!\x00\x16\x00\x0f\x00\xdc\x00\x8d\x00r\x00\x1b\x00X\x009\x00\x0e\x00g\x00\x14\x00%\x00\xea\x00\xf3\x00\x10\x00Q\x00\x06\x00\xbf\x00L\x00\xbd\x00b\x00\xcb\x00\xc8\x00i\x00\xfe\x00\x17\x00\x84\x00U\x00\xda\x00\xa3\x00\x80\x00\x81\x00\xf6\x00o\x00\xbc\x00\xed\x00R\x00{\x008\x00\x99\x00\xee\x00\xc7\x00\xf4\x00\x85\x00\xca\x00S\x00\xf0\x00\xb1\x00\xe6\x00\x1f\x00,\x00\x1d\x00B\x00+\x00\xa8\x00\xc9\x00\xde\x00w\x00d\x00\xb5\x00\xba\x00\x03\x00`\x00\xe1\x00\xd6\x00\xcf\x00\x9c\x00M\x002\x00\xdb\x00\x18\x00\xf9\x00\xce\x00\'\x00\xd4\x00\xe5\x00\xaa\x00\xb3\x00\xd0\xab\x11\x05\xc6\xa9\x7f\x01\x0c\x90}\x01"\x00\x8b\x00\x88\x00)\x00\xbe\x00\xd7\x00D\x00\x15\x00\x9a\x00c\x00@\x00A\x00\xb6\x00/\x00|\x00\xad\x00\x12\x00;\x00\xf8\x00Y\x00\xae\x00\x87\x00\xb4\x00E\x00\x8a\x00\x13\x00\xb0\x00q\x00\xa6\x00\xdf\x00\xec\x00\xdd\x00\x02\x00\xeb\x00h\x00\x89\x00\x9e\x007\x00$\x00u\x00z\x00\xc3\x00 \x00\xa1\x00\x96\x00\x8f\x00\\\x00\r\x00\xf2\x00\x9b\x00\xd8\x00\xb9\x00\x8e\x00\xe7\x00\x94\x00\xa5\x00j\x00s\x00\x90\x00\xd1\x00\x86\x00?\x00\xcc\x00=\x00\xe2\x00K\x00H\x00\xe9\x00~\x00\x97\x00\x04\x00\xd5\x00Z\x00#\x00\x00\x00\x01\x00v\x00\xef\x00<\x00m\x00\xd2\x00\xfb\x00\xb8\x00\x19\x00n\x00G\x00t\x00\x05\x00J\x00\xd3\x00p\x001\x00f\x00\x9f\x00\xac\x00\x9d\x00\xc2\x00\xab\x00(\x00I\x00^\x00\xf7\x00\xe4\x005\x00:\x00\x83\x00\xe0\x00a\x00V\x00O\x00\x1c\x00\xcd\x00\xb2\x00[\x00\x98\x00y\x00N\x00\xa7\x00T\x00e\x00*\x003\x00P\x00\x91\x00F\x00\xff\x00\x8c\xac\xfd\x00\xa2\xa9\x0b\x01\x08\xff\xa9\x01>\x00W\x00\xc4\x00\x95\x00\x1a\x00\xe3\x00\xc0\x00\xc1\x006\x00\xaf\x00\xfc\x00-\x00\x92\x00\xbb\x00x\x00\xd9\x00.\x00\x07\x004\x00\xc5\x00\n\x00\x93\x000\x00\xf1\x00&\x00_\x00l\x00]\x00\x82\x00k\x00\xe8\x00\t\x00\x1e\x00\xb7\x00\xa4\x00\xf5\x00\xfa\x00C\x00\xa0\x00!\x00\x16\x00\x0f\x00\xdc\x00\x8d\x00r\x00\x1b\x00X\x009\x00\x0e\x00g\x00\x14\x00%\x00\xea\x00\xf3\x00\x10\x00Q\x00\x06\x00\xbf\x00L\x00\xbd\x00b\x00\xcb\x00\xc8\x00i\x00\xfe\x00\x17\x00\x84\x00U\x00\xda\x00\xa3\x00\x80\x00\x81\x00\xf6\x00o\x00\xbc\x00\xed\x00R\x00{\x008\x00\x99\x00\xee\x00\xc7\x00\xf4\x00\x85\x00\xca\x00S\x00\xf0\x00\xb1\x00\xe6\x00\x1f\x00,\x00\x1d\x00B\x00+\x00\xa8\x00\xc9\x00\xde\x00w\x00d\x00\xb5\x00\xba\x00\x03\x00`\x00\xe1\x00\xd6\x00\xcf\x00\x9c\x00M\x002\x00\xdb\x00\x18\x00\xf9\x00\xce\x00\'\x00\xd4\x01\xe5\x01\xaa\x00\xb3\x00\xd0\x00\x11\x00\xc6\x00\x7f\x00\x0c\x00}\x00"\x00\x8b\x00\x88\xad)\x00\xbe\x00\xd7\x01D\xa9\x15\x01\x9a\x00c\x00@\x00A\x00\xb6\x00/\x00|\x00\xad\x00\x12\x00;\x00\xf8\x00Y\x00\xae\x00\x87\x00\xb4\x00E\x00\x8a\x00\x13\x00\xb0\x00q\x00\xa6\x00\xdf\x00\xec\x00\xdd\x00\x02\x00\xeb\x00h\x00\x89\x00\x9e\x007\x00$\x00u\x00z\x00\xc3\x00 \x00\xa1\x00\x96\x00\x8f\x00\\\x00\r\x00\xf2\x00\x9b\x00\xd8\x00\xb9\x00\x8e\x00\xe7\x00\x94\x00\xa5\x00j\x00s\x00\x90\x00\xd1\x00\x86\x00?\x00\xcc\x00=\x00\xe2\x00K\x00H\x00\xe9\x00~\x00\x97\x00\x04\x00\xd5\x00Z\x00#\x00\x00\x00\x01\x00v\x00\xef\x00<\x00m\x00\xd2\x00\xfb\x00\xb8\x00\x19\x00n\x00G\x00t\x00\x05\x00J\x00\xd3\x00p\x001\x00f\x00\x9f\x00\xac\x00\x9d\x00\xc2\x00\xab\x00(\x00I\x00^\x00\xf7\x00\xe4\x005\x00:\x00\x83\x00\xe0\x00a\x00V\x00O\x00\x1c\x00\xcd\x00\xb2\x00[\x00\x98\x00y\x00N\x00\xa7\x00T\x00e\x00*\x003\x00P\x03\x91\x01F\x03\xff\x00\x8c\x00\xfd\x07\xa2\x00\x0b\x00\x08\x02\xa9\x07>oWn\xc4c\x95\x00\x1a\xf0\xe3\x00\xc0\x00\xc1\x006\x00\xafa\xfcp-\x00\x92a\xbbpx\x08\xd9l.s\x07V4r\xc5\x00\nM\x93C09\xf10&6_,l5]\x00\x82skf\xe8r\t\x02\x1e\x1b\xb7p\xa4-\xf5e\xfaoCr\xa0e!/\x16r\x0fa\xdcc\x8dsr.\x1bwX\x009t\x0eUgl\x14\x00%r\xeam\xf3:\x10/Q2\x06.\xbf.L.\xbd/bp\xcb\x00\xc8\xc3ip\xfed\x17\x00\x84\x0cUa\xdaa\xa3i\x80i\x81i\xf6so@\xbc\x00\xed\x00R\x00{\x008a\x99d\xeeo\xc7o\xf4e\x85s\xca@S8\xf0\x00\xb1\x00\xe6\x00\x1fv,d\x1doBo+e\xa8s\xc9@\xde\x00w\x00d\x00\xb5\x00\xbav\x03d`o\xe1u\xd6c\xcfi\x9cnM?2\x00\xdb\x00\x18\x00\xf9\x00\xcep\'g\xd4U\xe5l\xaa\x00\xb3o\xd0j\x11c\xc6E\x7fc\x0cd}n"\x00\x8b\x00\x88\x00)\x00\xbe\x00\xd7\x00D\t\x15\x00\x9a\x00c\x00@\x04A\x01\xb6\x00/r|g\xads\x12e;\x00\xf8\x00Y\x00\xae\x00\x87\x00\xb4\x02E\r\x8ad\x13f\xb0uqt\xa6o\xdft\xec\x00\xdd\x00\x02\x00\xeb\x00h\x00\x89\x00\x9e\x007\x00$\x00u\x00z\x00\xc3\x00 \x00\xa1\x00\x96\x00\x8f\x00\\\x00\r\x00\xf2\x00\x9b\x00\xd8\x00\xb9\x00\x8e\x00\xe7\x00\x94\x00\xa5\x00j\x00s\x00\x90\x00\xd1\x00\x86\x00?\x00\xcc\x00=\x00\xe2\x00K\x00H\x00\xe9\x00~\x00\x97\x00\x04\x00\xd5\x00Z\x00#\x00\x00\x00\x01\x00v\x00\xef\x00<\x00m\x00\xd2\x00\xfb\x00\xb8\x00\x19\x00n\x00G\x00t\x00\x05\x00J\x00\xd3\x00p\x001\x00f\x00\x9f\x00\xac\x00\x9d\x00\xc2\x00\xab\x00(\x00I\x00^\x00\xf7\x00\xe4\x005\x00:\x00\x83\x00\xe0\x00a\x00V\x00O\x00\x1c\x00\xcd\x00\xb2\x00[\x00\x98\x00y\x00N\x00\xa7\x00T\x00e\x00*\x003\x00P\x00\x91\x00F\x00\xff\x00\x8c\x00\xfd\x00\xa2\x00\x0b\x00\x08\x00\xa9\x00>\x00W\x00\xc4\x00\x95\x00\x1a\x00\xe3\x00\xc0\x00\xc1\x00' + shake
            self.write(reply)
            self.setState('confirmClientHandshake')
            
    def confirmClientHandshake(self):
        if len(self.handshake_data) >= 1536:
            self.handshaked = True
            self.decoder.process(self.handshake_data[1536:])
            self.handshake_data = ''
            self.setState('handleEvents')

    def handshake(self):
        "Process a NetConnection client"
        self.dropper.setConsumer(self.transport)
        self.setState('readClientHandshake')
        
    def _publish(self, data, stream):
        name = data[3]
        stream.publish(name)
            
    def _FCPublish(self, data, stream):
        name = data[3]
        log.msg("Stream %s is coming from Flash Media Encoder" % name)
        
    def _play(self, data, stream):
        name = data[3]
        stream.play(name)
        
    def _FCUnpublish(self, data, stream):
        pass

class NetConnectionClient(Base):
    
    def __init__(self):
        Base.__init__(self)
        self.remoteConnected = False
        self.remotePublishing = False
        self.remotePlaying = False
    
    def readServerHandshake(self):
        if len(self.handshake_data) >= 3073:
            self.transport.write(self.handshake_data[1:1537])
            self.handshaked = True
            self.decoder.process(self.handshake_data[3073:])
            self.handshake_data = ''
            self.app = self.factory.app or "pms"
            self.tcUrl = self.factory.tcUrl or "rtmp://%s:%s/pms" % (self.ip, self.port)
            obj = {'videoCodecs' : 252.0,
                'audioCodecs' : 1639.0,
                'flashVer' : 'MAC 9,0,115,0',
                'app' : self.factory.app,
                'videoFunction' : 1.0,
                'capabilites' : 15.0,
                'fpad' : 0,
                'tcUrl' : self.factory.tcUrl,
                'swfUrl' : self.factory.swfUrl,
                'pageUrl' : self.factory.pageUrl}
            log.msg(obj)
            self.invoke('connect', 1.0, obj, insertNone=False)
            self.setState('handleEvents')

    def handshake(self):
        "Process a server side NetConnection"
        shake = chr(0x03) + struct.pack('!I', int(self.factory.core.get_uptime() * 1000.0))
        shake += chr(0x00) * 1532
        self.write(shake)
        self.setState('readServerHandshake')


